#include <assert.h>
#include <elf.h>
#include <errno.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define IMAGE_FILE "./image"
#define ARGS "[--extended] [--vm] <bootblock> <executable-file> ..."

#define OS_SIZE_LOC 0x1fc

int count = 0;          //count files

/* structure to store command line options */
static struct {
    int vm;
    int extended;
} options;

/* prototypes of local functions */
static void create_image(int nfiles, char *files[]);
static void error(char *fmt, ...);
static void read_ehdr(Elf64_Ehdr * ehdr, FILE * fp);
static void read_phdr(Elf64_Phdr * phdr, FILE * fp, int ph,
                      Elf64_Ehdr ehdr);
static void write_segment(Elf64_Ehdr ehdr, Elf64_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first);
static void write_os_size(int nbytes, FILE * img);

int main(int argc, char **argv)
{
    char *progname = argv[0];

    /* process command line options */
    options.vm = 0;
    options.extended = 0;
    while ((argc > 1) && (argv[1][0] == '-') && (argv[1][1] == '-')) {
        char *option = &argv[1][2];

        if (strcmp(option, "vm") == 0) {
            options.vm = 1;
        } else if (strcmp(option, "extended") == 0) {
            options.extended = 1;
        } else {
            error("%s: invalid option\nusage: %s %s\n", progname,
                  progname, ARGS);
        }
        argc--;
        argv++;
    }
    if (options.vm == 1) {
        error("%s: option --vm not implemented\n", progname);
    }
    if (argc < 3) {
        /* at least 3 args (createimage bootblock kernel) */
        error("usage: %s %s\n", progname, ARGS);
    }
    create_image(argc - 1, argv + 1);
    return 0;
}

static void create_image(int nfiles, char *files[])
{
    int ph, nbytes = 0, first = 1;
    FILE *fp, *img;
    Elf64_Ehdr ehdr;
    Elf64_Phdr phdr;

    /* open the image file */
    //add code in task 3
    img = fopen(IMAGE_FILE,"wb+");                      //read image file
    if(img == NULL){
        printf("ERROR: CANNOT OPEN FILE IMAGE!!!\n");  //cannot open the file IMAGE_FILE
        exit(1);                                       //kill the process
    }

    /* for each input file */
   
    while (nfiles-- > 0) {

        /* open input file */
        //add code in task 3
        fp = fopen(*files,"rb+");      //read only
        if(fp == NULL){
            printf("ERROR: CANNOT OPNE THE TARGET FILE!!!\n");
            exit(1);                    //To deal with some exceptions
        }

        /* read ELF header */
        read_ehdr(&ehdr, fp);
        printf("0x%04lx: %s\n", ehdr.e_entry, *files);

        /* for each program header */
        for (ph = 0; ph < ehdr.e_phnum; ph++) {

            /* read program header */
            read_phdr(&phdr, fp, ph, ehdr);

            /* write segment to the image */
            write_segment(ehdr, phdr, fp, img, &nbytes, &first);
        }
        fclose(fp);
        files++;
        count++;
        if(count >= 2){
            write_os_size(nbytes, img);
        }
    }
    fclose(img);
}

static void read_ehdr(Elf64_Ehdr * ehdr, FILE * fp)
{
    fread(ehdr, sizeof(Elf64_Ehdr), 1, fp);       //read
}                                              

static void read_phdr(Elf64_Phdr * phdr, FILE * fp, int ph,
                      Elf64_Ehdr ehdr)
{   //reloc fp from the begin of the file
    fseek(fp, (long)(ehdr.e_phoff + ph*ehdr.e_phentsize), SEEK_SET);  
    fread(phdr, sizeof(Elf64_Phdr), 1, fp); //read file
    
}

static void write_segment(Elf64_Ehdr ehdr, Elf64_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first)
{
    
    
    int fan_num;                                    //The number of the fans
    fan_num = (phdr.p_memsz % 512 == 0)? phdr.p_memsz/512: phdr.p_memsz/512 + 1;

    int msize = fan_num*512;                        //The Bytes need mallocing
    char a[msize];
    for(int i = 0; i < msize; i++){
        a[i] = 0;
    }                        //fill a with 0

    fseek(fp, (long)phdr.p_offset, SEEK_SET);           //set fp pointer(read position)
    fread(a, (long)phdr.p_filesz, 1,fp);                //read file
    
    fseek(img, (long)(*nbytes), SEEK_SET);              //set img pointer(write position)
    fwrite(a, (long)phdr.p_memsz, 1, img);              //write image file

    *nbytes += msize;                                 //nbyte + write bytes
    *first ++;                                        //only++ I do not really know what *first can do
}

int off = 0;
static void write_os_size(int nbytes, FILE * img)
{
    unsigned int kernel_size = nbytes/512 - 1 - off;        //Bootloader take one block
    int offset = OS_SIZE_LOC - (count - 2)*2;
    fseek(img, (long)offset, SEEK_SET);                     //find where to write the size
    printf("Pointer: %ld\n",ftell(img));
    printf("Kernel_size: %d count: %d img:%p\n",kernel_size,count,img);
    char size[2] = {kernel_size & 0xff, (kernel_size & 0xff00)>>8};     //little endien
    fwrite(size, 2, 1, img);

    off += kernel_size;      //write kernel size
}

/* print an error message and exit */
static void error(char *fmt, ...)
{
    va_list args;

    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
    va_end(args);
    if (errno != 0) {
        perror(NULL);
    }
    exit(EXIT_FAILURE);
}
